//
// Created by ECNU on 2024/2/22.
//

#ifndef NANO_FAISS_HEAP_H
#define NANO_FAISS_HEAP_H

#include <nanofaiss/utils/ordered_key_value.h>
#include <nanofaiss/impl/FaissAssert.h>

#include <cstring>

namespace faiss {

template<class C>
inline void heap_push(
        size_t k,
        typename C::T *bh_val,
        typename C::TI *bh_ids,
        typename C::T val,
        typename C::TI id) {
    // we use 1-based index
    bh_val--; bh_ids--;
    size_t i = k, i_father = 0;
    while (i > 1) {
        i_father = i >> 1;
        if (!C::cmp2(val, bh_val[i_father], id, bh_ids[i_father])) {
            break;
        }
        bh_val[i] = bh_val[i_father];
        bh_ids[i] = bh_ids[i_father];
        i = i_father;
    }
    bh_val[i] = val;
    bh_ids[i] = id;
}


/// Initialize the heap, k0 is the size of val and ids
/// k is the size of heap
template<class C>
inline void heap_heapify(
        size_t k,
        typename C::T *bh_val,
        typename C::TI *bh_ids,
        const typename C::T *val = nullptr,
        const typename C::TI *ids = nullptr,
        size_t k0 = 0) {
    if (k0 > 0)
        FAISS_THROW_IF_NOT(val);

    if (ids) {
        // because we used 1-based index
        for (size_t i = 0; i != k0; i++)
            heap_push<C>(i + 1, bh_val, bh_ids, val[i], ids[i]);
    } else {
        for (size_t i = 0; i != k0; i++)
            heap_push<C>(i + 1, bh_val, bh_ids, val[i], i);
    }

    for (size_t i = k0; i < k; i++) {
        bh_val[i] = C::neural();
        bh_ids[i] = -1;
    }
}

/// pops the elements from the heap, but elements in k is undefined
template<class C>
inline void heap_pop(size_t k, typename C::T *bh_val, typename C::TI *bh_ids) {
    bh_val--; bh_ids--; // use 1-based index, then [1...k], but only use [1...k-1]
    typename C::T val = bh_val[k];
    typename C::TI id = bh_ids[k];
    size_t i = 1, c1, c2;
    while (i * 2 <= k) {
        c1 = i << 1;
        c2 = c1 + 1;
        // answer reserves in c1;
        if (c2 < k && C::cmp2(bh_val[c2], bh_val[c1], bh_ids[c2], bh_ids[c1]))
            c1 = c2;
        if (c1 >= k || C::cmp2(val, bh_val[c1], id, bh_ids[c1]))
            break;
        bh_val[i] = bh_val[c1];
        bh_ids[i] = bh_ids[c1];
        i = c1;
    }
    bh_val[i] = val;
    bh_ids[i] = id;
}

/// reorder the heap into the order vector
/// because the rest of the vector is -1
/// [Max, Max, Max, normal, normal]
template<class C>
inline size_t heap_reorder(
        size_t k,
        typename C::T *vals,
        typename C::TI *ids) {
    size_t i, ii;
    for (i = 0, ii = 0; i != k; i++) {
        typename C::T val = vals[0];
        typename C::TI id = ids[0];
        heap_pop<C>(k - i, vals, ids);
        vals[k - ii - 1] = val;
        ids[k - ii - 1] = id;
        if (id != -1) // the rest may be max or minimum
            ii++;
    }
    size_t nel = ii;

    // move the meaningful element to the front
    memmove(vals, vals + k - ii, ii * sizeof(*vals));
    memmove(ids, ids + k - ii, ii * sizeof(*ids));

    // padding the none-meaningful point
    for (; ii != k; ii++) {
        vals[ii] = C::neural();
        ids[ii] = -1;
    }
    return nel;
}

template<class C>
inline void heap_replace_top(
        size_t k,
        typename C::T *bh_val,
        typename C::TI *bh_ids,
        typename C::T val,
        typename C::TI id) {
    bh_val--; bh_ids--;
    size_t i = 1, c1, c2;
    while (i * 2 <= k) {
        c1 = i << 1;
        c2 = c1 + 1;
      if (c2 <= k && C::cmp2(bh_val[c2], bh_val[c1], bh_ids[c2], bh_ids[c1]))
          c1 = c2;
      if (c1 > k || C::cmp2(val, bh_val[c1], id, bh_ids[c1]))
          break;
      bh_val[i] = bh_val[c1];
      bh_ids[i] = bh_ids[c1];
      i = c1;
    }
    bh_val[i] = val;
    bh_ids[i] = id;
}


/** C -> CMax or CMin
 * Heap array is a multiple heap
 * */
template<typename C>
struct HeapArray {
    typedef typename C::T T;
    typedef typename C::TI TI;

    size_t nh; /// number of heap
    size_t k; /// size per heap
    TI *ids; /// id of vector (nh * k)
    T *val; /// distance (nh * k)

    ///
    T* get_val(size_t key) {
        return val + key * k;
    }

    TI* get_id(size_t key) {
        return ids + key * k;
    }


    /// prepare all heap before adding
    void heapify();


    /// reorder all the heap
    void reorder();


    /** add nj elements to heap i0:i0+nj
     * @param nj    heap number
     * @param vin   input ni * nj
     * @param j0    is the id need to add
     * @param i0    first heap to add
     * @param ni    update heap number, if -1 then equal to nh
     */
    void addn(
            size_t nj,
            const T *vin,
            TI j0 = 0,
            size_t i0 = 0,
            int64_t ni = -1);

    void add_with_ids(
            size_t nj,
            const T *vin,
            const TI *id_in = nullptr,
            int64_t id_stride = 0,
            size_t i0 = 0,
            int64_t ni = -1);


};

template<typename T, typename TI>
struct QueryParam {
    size_t nq;
    size_t k;
    T *distances;
    TI *ids;
};





/** Define useful Heap */
typedef HeapArray<CMin<float, int64_t>> float_minheap_array_t;
typedef HeapArray<CMin<int, int64_t>> int_minheap_array_t;

typedef HeapArray<CMax<float, int64_t>> float_maxheap_array_t;
typedef HeapArray<CMax<int, int64_t>> int_maxheap_array_t;

/** Define useful QueryParam */
typedef QueryParam<float, int64_t> float_query_t;

}
#endif // NANO_FAISS_HEAP_H
